package com.wxl.asmgradleplugin.myplugin

import javassist.ClassPool
import javassist.CtClass
import org.gradle.api.DefaultTask
import org.gradle.api.file.ConfigurableFileCollection
import org.gradle.api.file.Directory
import org.gradle.api.file.DirectoryProperty
import org.gradle.api.file.FileCollection
import org.gradle.api.file.FileType
import org.gradle.api.file.RegularFile
import org.gradle.api.file.RegularFileProperty
import org.gradle.api.provider.ListProperty
import org.gradle.api.provider.Property
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.InputDirectory
import org.gradle.api.tasks.InputFiles
import org.gradle.api.tasks.Optional
import org.gradle.api.tasks.OutputDirectory
import org.gradle.api.tasks.OutputFile
import org.gradle.api.tasks.PathSensitive
import org.gradle.api.tasks.PathSensitivity
import org.gradle.api.tasks.TaskAction
import org.gradle.work.ChangeType
import org.gradle.work.FileChange
import org.gradle.work.Incremental
import org.gradle.work.InputChanges
import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassWriter
import org.objectweb.asm.Opcodes
import java.io.File
import java.io.InputStream
import java.util.jar.JarEntry
import java.util.jar.JarFile
import java.util.jar.JarOutputStream


abstract class MyTransformTask : DefaultTask() {

    companion object {
        private const val DOT_CLASS = ".class"
    }

    @get:Input
    abstract val androidSdkDir: Property<String>

    @get:Input
    abstract val compileClasspath: Property<FileCollection>

    @get:InputFiles
    @get:PathSensitive(PathSensitivity.RELATIVE)
    abstract val allJars: ListProperty<RegularFile>

    //    @get:Incremental
    @get:InputFiles
    @get:PathSensitive(PathSensitivity.RELATIVE)
    abstract val allDirs: ListProperty<Directory>

//    @get:InputFile
//    abstract val mergedManifest: RegularFileProperty

//    @get:InputFile
//    abstract val producedManifest: RegularFileProperty

    @get:OutputFile
    abstract val output: RegularFileProperty

//    @get:Incremental
//    @get:InputFiles
//    @get:PathSensitive(PathSensitivity.RELATIVE)
//    val inputDirs: ConfigurableFileCollection =
//        project.objects.fileCollection().from(allDirs)

//    @get:Optional
//    @get:Incremental
//    @get:InputDirectory
//    abstract val inputDir: DirectoryProperty
//
//    @get:Optional
//    @get:Incremental
//    @get:InputDirectory
//    abstract val inputDir2: DirectoryProperty
//
//    @get:Optional
//    @get:OutputDirectory
//    abstract val outputDir: DirectoryProperty
//

    private val classPool = ClassPool.getDefault()
    private val jarPaths = mutableSetOf<String>()
    private val javassist = false

    @TaskAction
    fun taskAction(changes: InputChanges) {

        println("=== MyTransformTask changes $changes")

        val androidSdkDir = androidSdkDir.get()
        println("=== androidSdkDir: $androidSdkDir")
        println("=== compileClasspath: $compileClasspath")
        classPool.appendClassPath(androidSdkDir)
        compileClasspath.get().forEach {
            classPool.appendClassPath(it.absolutePath)
        }

//        println("=== mergedManifest ${mergedManifest.get().asFile.absolutePath}")
//        println("=== producedManifest ${producedManifest.get().asFile.absolutePath}")
        if (changes.isIncremental) {
            println("------ MyTransformTask changes isIncremental")
//            handleChanges(changes, inputDir)
//            handleChanges(changes, inputDir2)
//            handleChanges(changes, inputDirs)
//            return
        }
        val outputFile = output.get().asFile
        println("=== outputFile ${outputFile.absolutePath}")
        JarOutputStream(outputFile.outputStream()).use { jarOutput ->
            processJars(jarOutput)
            processDirs(jarOutput)
        }

        println("=== MyTransformTask execute done")
    }

    private fun processDirs(jarOutput: JarOutputStream) {
        allDirs.get().forEach { directory ->
            val dir = directory.asFile
            println("=== processDir dir: ${dir.absolutePath}")
            dir.walk().filter { it.isFile && it.extension == "class" }.forEach { file ->
                processFile(file, dir, jarOutput)
            }
        }
    }

    private fun processFile(file: File, dir: File, jarOutput: JarOutputStream) {
        //                println("=== processDir file: ${file.absolutePath}")
        // replace \ to /
        val classPath = file.relativeTo(dir).path.replace(File.separatorChar, '/')
        //                println("=== processDir classPath: $classPath")
        var bytecode: ByteArray? = null
        if (classPath.startsWith(Consts.PK)) {
            println("=== processDir startsWith file: ${file.absolutePath}")
            println("=== processDir startsWith classPath: $classPath")
            bytecode = insertCode(classPath, file.inputStream())
        }
        if (bytecode != null) {
            jarOutput.writeEntity(classPath, bytecode)
        } else {
            jarOutput.writeEntity(classPath, file.inputStream())
        }
    }


    /**
     * 1. 使用jarOutput.putNextEntry创建一个和输入jarEntry同名的新入口到输出jar中。
     * 2. 使用jarFile.getInputStream获取输入jarEntry对应的输入流。
     * 3. 使用InputStream的copyTo方法将输入流复制到JarOutputStream中,完成内容的复制。
     * 4. 用use方法保证流被关闭。
     * 这样,就实现了从输入的JarFile选择性地复制jarEntry到输出的jar文件中。
     */
    private fun processJars(jarOutput: JarOutputStream) {
        allJars.get().forEach { file ->
            JarFile(file.asFile).use { jarFile ->
//                println("=== processJar: jarFile name=${jarFile.name} size=${jarFile.size()}")
                jarFile.entries().iterator().forEach { jarEntry ->
                    val entryName = jarEntry.name
//                    println("=== processJar: entryName= $entryName")
                    val inputStream = jarFile.getInputStream(jarEntry)
                    var bytecode: ByteArray? = null
                    if (entryName.endsWith(DOT_CLASS) && entryName.startsWith(Consts.PK)) {
                        println("=== processJar: startsWith jarFile name=${jarFile.name} size=${jarFile.size()}")
                        println("=== processJar: startsWith entryName= $entryName")
                        bytecode = insertCode(jarEntry.name, inputStream)
                    }
                    if (bytecode != null) {
                        jarOutput.writeEntity(entryName, bytecode)
                    } else {
                        jarOutput.writeEntity(entryName, inputStream)
                    }
                }
            }
        }
    }

    private fun JarOutputStream.writeEntity(name: String, inputStream: InputStream) {
        if (jarPaths.contains(name)) {
            printDuplicatedMessage(name)
        } else {
            runCatching {
                putNextEntry(JarEntry(name))
                inputStream.use {
                    it.copyTo(this)
                }
            }.onFailure { e ->
                println("=== === Copy jar entry failed. [entry:$name] e=$e")
            }
            closeEntry()
            jarPaths.add(name)
        }
    }

    private fun JarOutputStream.writeEntity(relativePath: String, bytecode: ByteArray) {
        if (jarPaths.contains(relativePath)) {
            printDuplicatedMessage(relativePath)
        } else {
            putNextEntry(JarEntry(relativePath))
            write(bytecode)
            closeEntry()
            jarPaths.add(relativePath)
        }
    }

    private fun printDuplicatedMessage(name: String) {
//        println("=== Cannot add ${name}, because output Jar already has file with the same name.")
    }

    private fun insertCode(name: String, inputStream: InputStream): ByteArray? {
        return if (javassist) {
            insertCodeJavassist(name, inputStream)
        } else {
            insertCodeASM(inputStream)
        }
    }

    private fun insertCodeASM(inputStream: InputStream): ByteArray? {
        try {
            val cr = ClassReader(inputStream)
            val cw = ClassWriter(cr, ClassWriter.COMPUTE_FRAMES)
            cr.accept(LogClassVisitor(Opcodes.ASM9, cw), ClassReader.EXPAND_FRAMES)
            return cw.toByteArray()
        } catch (ex: Throwable) {
            println("=== === insertCodeASM: $ex")
        }
        return null
    }

    private fun insertCodeJavassist(name: String, inputStream: InputStream): ByteArray? {
        try {
//            val className = name.removeSuffix(DOT_CLASS).replace('/', '.')
//            println("=== insertCodeJavassist className: $className")
            val ctClass = inputStream.use {
                classPool.makeClass(it)
            }
//            val ctClass = classPool.getOrNull(className)
//                ?: inputStream.use {
//                    println("=== ctClass.makeClass ")
//                    classPool.makeClass(it)
//                }
//        val ctClass = inputStream.use {
//            val classPool = ClassPool()
//            classPool.appendSystemPath()
//            classPool.makeClass(it)
//        }
            if (ctClass.isFrozen) {
                println("=== === ctClass.isFrozen ")
                ctClass.defrost()
            }
            try {
                val method = ctClass.getDeclaredMethod("test")
                method.insertBefore("""System.out.println(">>> Injected by Javassist <<<");""")
            } catch (e: Exception) {
                println("=== insertCodeJavassist: e=$e")
            }
            try {
                if (ctClass.superclass?.name == "androidx.activity.ComponentActivity") {
                    println("=== insertCodeJavassist ctClass: $ctClass")
                    val params = arrayOf<CtClass>(classPool.get("android.os.Bundle"))
                    val onCreateCtMethod = ctClass.getDeclaredMethod("onCreate", params)
                    onCreateCtMethod.insertBefore("android.util.Log.i(\"=== Javassist ===\", \"=== onCreate\");")
                }
            } catch (e: Exception) {
                println("=== insertCodeJavassist: e=$e")
            }
            val bytecode = ctClass.toBytecode()
            ctClass.detach()
            return bytecode
        } catch (tr: Throwable) {
            println("=== === insertCodeJavassist: tr=$tr")
        }
        return null
    }

    private fun handleChanges(changes: InputChanges, changeDir: DirectoryProperty) {
        changes.getFileChanges(changeDir).forEach { change ->
            handleChange(change)
        }
    }

    private fun handleChanges(changes: InputChanges, changeDir: FileCollection) {
        changes.getFileChanges(changeDir).forEach { change ->
            handleChange(change)
        }
    }

    private fun handleChange(change: FileChange) {
        if (change.fileType == FileType.FILE) {
            val file = change.file
            println("=== change ${change.changeType} $file")
            //                val targetFile = outputDir.file(change.normalizedPath).get().asFile
            when (change.changeType) {
                ChangeType.REMOVED -> {
                    println("=== execute REMOVED")
                }

                ChangeType.ADDED -> {
                    println("=== execute ADDED")
                    allDirs.get().forEach { directory ->
                        val dir = directory.asFile
                        if (file.absolutePath.startsWith(dir.absolutePath)) {
                            JarOutputStream(output.get().asFile.outputStream()).use { jarOutput ->
                                processFile(file, dir, jarOutput)
                            }
                        }
                    }
                }

                else -> {
                    println("=== execute other")
                }
            }
        }
    }
}